// Copyright 2022 jeff.li. and/or its affiliates.
// Acknowledgement:
// Taken from https://github.com/pytorch/pytorch/blob/release/1.11/c10/util/Half.h
// with fixes applied:
// - change namespace to tbir::runtime for fix conflict with pytorch

#pragma once

#include <cstring>
#include <limits>

#include <tbir/runtime/runtime_port.h>

#ifdef __CUDACC__
#include <cuda_fp16.h>
#endif

#ifdef __HIPCC__
#include <hip/hip_fp16.h>
#endif

#ifdef __SYCL_DEVICE_ONLY__
#include <CL/sycl.hpp>
#endif

TBIR_CLANG_DIAGNOSTIC_PUSH()
#if TBIR_CLANG_HAS_WARNING("-Wimplicit-int-float-conversion")
TBIR_CLANG_DIAGNOSTIC_IGNORE("-Wimplicit-int-float-conversion")
#endif

namespace tbir::runtime {

/// Constructors

    inline TBIR_RUNTIME_HOST_DEVICE Half::Half(float value) {
#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__)
        x = __half_as_short(__float2half(value));
#elif defined(__SYCL_DEVICE_ONLY__)
        x = sycl::bit_cast<uint16_t>(sycl::half(value));
#else
        x = detail::fp16_ieee_from_fp32_value(value);
#endif
    }

/// Implicit conversions

    inline TBIR_RUNTIME_HOST_DEVICE Half::operator float() const {
#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__)
        return __half2float(*reinterpret_cast<const __half*>(&x));
#elif defined(__SYCL_DEVICE_ONLY__)
        return float(sycl::bit_cast<sycl::half>(x));
#else
        return detail::fp16_ieee_to_fp32_value(x);
#endif
    }

#if defined(__CUDACC__) || defined(__HIPCC__)
    inline TBIR_RUNTIME_HOST_DEVICE Half::Half(const __half& value) {
      x = *reinterpret_cast<const unsigned short*>(&value);
    }
    inline TBIR_RUNTIME_HOST_DEVICE Half::operator __half() const {
      return *reinterpret_cast<const __half*>(&x);
    }
#endif

// CUDA intrinsics

#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 350)) || (defined(__clang__) && defined(__CUDA__))
    inline __device__ Half __ldg(const Half* ptr) {
      return __ldg(reinterpret_cast<const __half*>(ptr));
    }
#endif

/// Arithmetic

    inline TBIR_RUNTIME_HOST_DEVICE Half operator+(const Half &a, const Half &b) {
        return static_cast<float>(a) + static_cast<float>(b);
    }

    inline TBIR_RUNTIME_HOST_DEVICE Half operator-(const Half &a, const Half &b) {
        return static_cast<float>(a) - static_cast<float>(b);
    }

    inline TBIR_RUNTIME_HOST_DEVICE Half operator*(const Half &a, const Half &b) {
        return static_cast<float>(a) * static_cast<float>(b);
    }

    inline TBIR_RUNTIME_HOST_DEVICE Half operator/(const Half &a, const Half &b)
    __ubsan_ignore_float_divide_by_zero__ {
        return static_cast<float>(a) / static_cast<float>(b);
    }

    inline TBIR_RUNTIME_HOST_DEVICE Half operator-(const Half &a) {
#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 530) || defined(__HIP_DEVICE_COMPILE__)
        return __hneg(a);
#else
        return -static_cast<float>(a);
#endif
    }

    inline TBIR_RUNTIME_HOST_DEVICE Half &operator+=(Half &a, const Half &b) {
        a = a + b;
        return a;
    }

    inline TBIR_RUNTIME_HOST_DEVICE Half &operator-=(Half &a, const Half &b) {
        a = a - b;
        return a;
    }

    inline TBIR_RUNTIME_HOST_DEVICE Half &operator*=(Half &a, const Half &b) {
        a = a * b;
        return a;
    }

    inline TBIR_RUNTIME_HOST_DEVICE Half &operator/=(Half &a, const Half &b) {
        a = a / b;
        return a;
    }

/// Arithmetic with floats

    inline TBIR_RUNTIME_HOST_DEVICE float operator+(Half a, float b) {
        return static_cast<float>(a) + b;
    }

    inline TBIR_RUNTIME_HOST_DEVICE float operator-(Half a, float b) {
        return static_cast<float>(a) - b;
    }

    inline TBIR_RUNTIME_HOST_DEVICE float operator*(Half a, float b) {
        return static_cast<float>(a) * b;
    }

    inline TBIR_RUNTIME_HOST_DEVICE float operator/(Half a, float b)
    __ubsan_ignore_float_divide_by_zero__ {
        return static_cast<float>(a) / b;
    }

    inline TBIR_RUNTIME_HOST_DEVICE float operator+(float a, Half b) {
        return a + static_cast<float>(b);
    }

    inline TBIR_RUNTIME_HOST_DEVICE float operator-(float a, Half b) {
        return a - static_cast<float>(b);
    }

    inline TBIR_RUNTIME_HOST_DEVICE float operator*(float a, Half b) {
        return a * static_cast<float>(b);
    }

    inline TBIR_RUNTIME_HOST_DEVICE float operator/(float a, Half b)
    __ubsan_ignore_float_divide_by_zero__ {
        return a / static_cast<float>(b);
    }

    inline TBIR_RUNTIME_HOST_DEVICE float &operator+=(float &a, const Half &b) {
        return a += static_cast<float>(b);
    }

    inline TBIR_RUNTIME_HOST_DEVICE float &operator-=(float &a, const Half &b) {
        return a -= static_cast<float>(b);
    }

    inline TBIR_RUNTIME_HOST_DEVICE float &operator*=(float &a, const Half &b) {
        return a *= static_cast<float>(b);
    }

    inline TBIR_RUNTIME_HOST_DEVICE float &operator/=(float &a, const Half &b) {
        return a /= static_cast<float>(b);
    }

    /// Arithmetic with doubles

    inline TBIR_RUNTIME_HOST_DEVICE double operator+(Half a, double b) {
        return static_cast<double>(a) + b;
    }

    inline TBIR_RUNTIME_HOST_DEVICE double operator-(Half a, double b) {
        return static_cast<double>(a) - b;
    }

    inline TBIR_RUNTIME_HOST_DEVICE double operator*(Half a, double b) {
        return static_cast<double>(a) * b;
    }

    inline TBIR_RUNTIME_HOST_DEVICE double operator/(Half a, double b)
    __ubsan_ignore_float_divide_by_zero__ {
        return static_cast<double>(a) / b;
    }

    inline TBIR_RUNTIME_HOST_DEVICE double operator+(double a, Half b) {
        return a + static_cast<double>(b);
    }

    inline TBIR_RUNTIME_HOST_DEVICE double operator-(double a, Half b) {
        return a - static_cast<double>(b);
    }

    inline TBIR_RUNTIME_HOST_DEVICE double operator*(double a, Half b) {
        return a * static_cast<double>(b);
    }

    inline TBIR_RUNTIME_HOST_DEVICE double operator/(double a, Half b)
    __ubsan_ignore_float_divide_by_zero__ {
        return a / static_cast<double>(b);
    }

    /// Arithmetic with ints

    inline TBIR_RUNTIME_HOST_DEVICE Half operator+(Half a, int b) {
        return a + static_cast<Half>(b);
    }

    inline TBIR_RUNTIME_HOST_DEVICE Half operator-(Half a, int b) {
        return a - static_cast<Half>(b);
    }

    inline TBIR_RUNTIME_HOST_DEVICE Half operator*(Half a, int b) {
        return a * static_cast<Half>(b);
    }

    inline TBIR_RUNTIME_HOST_DEVICE Half operator/(Half a, int b) {
        return a / static_cast<Half>(b);
    }

    inline TBIR_RUNTIME_HOST_DEVICE Half operator+(int a, Half b) {
        return static_cast<Half>(a) + b;
    }

    inline TBIR_RUNTIME_HOST_DEVICE Half operator-(int a, Half b) {
        return static_cast<Half>(a) - b;
    }

    inline TBIR_RUNTIME_HOST_DEVICE Half operator*(int a, Half b) {
        return static_cast<Half>(a) * b;
    }

    inline TBIR_RUNTIME_HOST_DEVICE Half operator/(int a, Half b) {
        return static_cast<Half>(a) / b;
    }

    //// Arithmetic with int64_t

    inline TBIR_RUNTIME_HOST_DEVICE Half operator+(Half a, int64_t b) {
        return a + static_cast<Half>(b);
    }

    inline TBIR_RUNTIME_HOST_DEVICE Half operator-(Half a, int64_t b) {
        return a - static_cast<Half>(b);
    }

    inline TBIR_RUNTIME_HOST_DEVICE Half operator*(Half a, int64_t b) {
        return a * static_cast<Half>(b);
    }

    inline TBIR_RUNTIME_HOST_DEVICE Half operator/(Half a, int64_t b) {
        return a / static_cast<Half>(b);
    }

    inline TBIR_RUNTIME_HOST_DEVICE Half operator+(int64_t a, Half b) {
        return static_cast<Half>(a) + b;
    }

    inline TBIR_RUNTIME_HOST_DEVICE Half operator-(int64_t a, Half b) {
        return static_cast<Half>(a) - b;
    }

    inline TBIR_RUNTIME_HOST_DEVICE Half operator*(int64_t a, Half b) {
        return static_cast<Half>(a) * b;
    }

    inline TBIR_RUNTIME_HOST_DEVICE Half operator/(int64_t a, Half b) {
        return static_cast<Half>(a) / b;
    }

    /// NOTE: we do not define comparisons directly and instead rely on the implicit
    /// conversion from c10::Half to float.

}  // namespace tbir::runtime

namespace std {

    template<>
    class numeric_limits<::tbir::runtime::Half> {
    public:
        static constexpr bool is_specialized = true;
        static constexpr bool is_signed = true;
        static constexpr bool is_integer = false;
        static constexpr bool is_exact = false;
        static constexpr bool has_infinity = true;
        static constexpr bool has_quiet_NaN = true;
        static constexpr bool has_signaling_NaN = true;
        static constexpr auto has_denorm = numeric_limits<float>::has_denorm;
        static constexpr auto has_denorm_loss = numeric_limits<float>::has_denorm_loss;
        static constexpr auto round_style = numeric_limits<float>::round_style;
        static constexpr bool is_iec559 = true;
        static constexpr bool is_bounded = true;
        static constexpr bool is_modulo = false;
        static constexpr int digits = 11;
        static constexpr int digits10 = 3;
        static constexpr int max_digits10 = 5;
        static constexpr int radix = 2;
        static constexpr int min_exponent = -13;
        static constexpr int min_exponent10 = -4;
        static constexpr int max_exponent = 16;
        static constexpr int max_exponent10 = 4;
        static constexpr auto traps = numeric_limits<float>::traps;
        static constexpr auto tinyness_before = numeric_limits<float>::tinyness_before;

        static constexpr ::tbir::runtime::Half min() {
            return ::tbir::runtime::Half(0x0400, ::tbir::runtime::Half::from_bits());
        }

        static constexpr ::tbir::runtime::Half lowest() {
            return ::tbir::runtime::Half(0xFBFF, ::tbir::runtime::Half::from_bits());
        }

        static constexpr ::tbir::runtime::Half max() {
            return ::tbir::runtime::Half(0x7BFF, ::tbir::runtime::Half::from_bits());
        }

        static constexpr ::tbir::runtime::Half epsilon() {
            return ::tbir::runtime::Half(0x1400, ::tbir::runtime::Half::from_bits());
        }

        static constexpr ::tbir::runtime::Half round_error() {
            return ::tbir::runtime::Half(0x3800, ::tbir::runtime::Half::from_bits());
        }

        static constexpr ::tbir::runtime::Half infinity() {
            return ::tbir::runtime::Half(0x7C00, ::tbir::runtime::Half::from_bits());
        }

        static constexpr ::tbir::runtime::Half quiet_NaN() {
            return ::tbir::runtime::Half(0x7E00, ::tbir::runtime::Half::from_bits());
        }

        static constexpr ::tbir::runtime::Half signaling_NaN() {
            return ::tbir::runtime::Half(0x7D00, ::tbir::runtime::Half::from_bits());
        }

        static constexpr ::tbir::runtime::Half denorm_min() {
            return ::tbir::runtime::Half(0x0001, ::tbir::runtime::Half::from_bits());
        }
    };

}  // namespace std

TBIR_CLANG_DIAGNOSTIC_POP()
